import os
import numpy as np
import pybullet as p
import pybullet_data
import time

import math
class WorldCreation:


    def __init__(self, pid, robot_type='jaco', time_step=0.02, np_random=None):
        self.id = pid
        self.robot_type = robot_type
        self.time_step = time_step
        self.np_random = np_random
        self.directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'assets')

    def create_new_world(self, print_joints=False):

        p.resetSimulation(physicsClientId=self.id)
        # Configure camera position
        # p.resetDebugVisualizerCamera(cameraDistance=1.75, cameraYaw=-25, cameraPitch=-45, cameraTargetPosition=[-0.2, 0, 0.4], physicsClientId=self.id)
        # p.configureDebugVisualizer(p.COV_ENABLE_MOUSE_PICKING, 0, physicsClientId=self.id)
        # p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.id)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        plane_id = p.loadURDF("plane.urdf", useMaximalCoordinates=True)
        # Disable rendering during creation
        # p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 0, physicsClientId=self.id)

        p.setTimeStep(self.time_step, physicsClientId=self.id)
        # Disable real time simulation so that the simulation only advances when we call stepSimulation
        # p.setRealTimeSimulation(0, physicsClientId=self.id)

        # Jaco
        self.init_jaco(print_joints)
        


    def init_jaco(self, print_joints=False):
        print("helloworld")
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        robot_id = p.loadURDF('r2d2.urdf',basePosition=[0, 0, 0.05], useMaximalCoordinates=True,physicsClientId=self.id)

        self.new_robot(robot_id)




    def set_gripper_open_position(self, robot, position=0, left=True, set_instantly=False, indices=None):
        indices_new = [9, 11, 13]
        positions = [position, position, position]
        if indices is None:
            indices = indices_new

        if set_instantly:
            for i, j in enumerate(indices):
                p.resetJointState(robot, jointIndex=j, targetValue=positions[i], targetVelocity=0, physicsClientId=self.id)
        p.setJointMotorControlArray(robot, jointIndices=indices, controlMode=p.POSITION_CONTROL, targetPositions=positions, positionGains=np.array([0.05]*len(indices)), forces=[500]*len(indices), physicsClientId=self.id)


    def setup_robot_joints(self, robot, robot_joint_indices, lower_limits, upper_limits, randomize_joint_positions=False, default_positions=[1, 1, 0, -1.75, 0, -1.1, -0.5], tool=None):

       
        default_positions[default_positions < lower_limits] = lower_limits[default_positions < lower_limits]
        default_positions[default_positions > upper_limits] = upper_limits[default_positions > upper_limits]
        for i in range(len(robot_joint_indices)):
            p.resetJointState(robot, jointIndex=i, targetValue=default_positions[i], targetVelocity=0, physicsClientId=self.id)


    def print_joint_info(self, body, show_fixed=True):
        joint_names = []
        for j in range(p.getNumJoints(body, physicsClientId=self.id)):
            if show_fixed or p.getJointInfo(body, j, physicsClientId=self.id)[2] != p.JOINT_FIXED:
                print(p.getJointInfo(body, j, physicsClientId=self.id))
                joint_names.append((j, p.getJointInfo(body, j, physicsClientId=self.id)[1]))
        print(joint_names)


    def enforce_joint_limits(self, body):
        # Enforce joint limits
        joint_states = p.getJointStates(body, jointIndices=list(range(p.getNumJoints(body, physicsClientId=self.id))), physicsClientId=self.id)
        joint_positions = np.array([x[0] for x in joint_states])
        lower_limits = []
        upper_limits = []
        for j in [2,3,4,5,6,7,9,10,11,12,13,14]:
            joint_info = p.getJointInfo(body, j, physicsClientId=self.id)
            joint_name = joint_info[1]
            joint_pos = joint_positions[j]
            lower_limit = joint_info[8]
            upper_limit = joint_info[9]
            if lower_limit == 0 and upper_limit == -1:
                lower_limit = -1e10
                upper_limit = 1e10
            lower_limits.append(lower_limit)
            upper_limits.append(upper_limit)
            # print(joint_name, joint_pos, lower_limit, upper_limit)
            if joint_pos < lower_limit:
                p.resetJointState(body, jointIndex=j, targetValue=lower_limit, targetVelocity=0, physicsClientId=self.id)
            elif joint_pos > upper_limit:
                p.resetJointState(body, jointIndex=j, targetValue=upper_limit, targetVelocity=0, physicsClientId=self.id)
        lower_limits = np.array(lower_limits)
        upper_limits = np.array(upper_limits)
        return lower_limits, upper_limits

    def new_robot(self,robot_id):
        p.setGravity(0, 0, -9.8)
        p.setRealTimeSimulation(1)
        p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1)
        # 如果是轮子的关节，则为马达配置参数，否则禁用马达
        for i in range(p.getNumJoints(robot_id)):
            if "wheel" in p.getJointInfo(robot_id, i)[1].decode("utf-8"):           
                p.setJointMotorControl2(
                    bodyUniqueId=robot_id,
                    jointIndex=i,
                    controlMode=p.VELOCITY_CONTROL,
                    targetVelocity=5,
                    force=60)
            else:
                p.setJointMotorControl2(
                    bodyUniqueId=robot_id,
                    jointIndex=i,
                    controlMode=p.VELOCITY_CONTROL,
                    force=0)

        # 如果不需要将激光可视化出来，置为False
        useDebugLine = True 
        hitRayColor = [0, 1, 0]
        missRayColor = [1, 0, 0]

        rayLength = 10          # 激光长度
        rayNum = 360             # 激光数量


        while True:
            p.stepSimulation()
            # 获取需要发射得rayNum激光的起点与终点
            begins, _ = p.getBasePositionAndOrientation(robot_id)


            rayFroms = [
                [
                    begins[0],
                    begins[1],
                    begins[2]+0.2
                ] 
            for i in range(rayNum)]
            rayTos = [
                [
                    begins[0] + rayLength * math.cos(2 * math.pi * float(i) / rayNum),
                    begins[1] + rayLength * math.sin(2 * math.pi * float(i) / rayNum),
                    begins[2]+0.3
                ] 
            for i in range(rayNum)]

            # 调用激光探测函数
            results = p.rayTestBatch(rayFroms, rayTos)

            # 染色前清楚标记
            p.removeAllUserDebugItems()

            # 根据results结果给激光染色
            for index, result in enumerate(results):
                if result[0] == -1:
                    p.addUserDebugLine(rayFroms[index], rayTos[index], missRayColor)
                else:
                    p.addUserDebugLine(rayFroms[index], rayTos[index], hitRayColor)



use_gui = True
if use_gui:
    serve_id = p.connect(p.GUI)
else:
    serve_id = p.connect(p.DIRECT)
x = WorldCreation(serve_id)
x.create_new_world()